"""
有向无环图(DAG)管理器
"""
from typing import Dict, List, Set, Any
import networkx as nx
from core.logger import get_logger

# 获取日志记录器
logger = get_logger(__name__)


class DAG:
    """有向无环图，用于表示Pipeline的执行流程"""

    def __init__(self):
        self.graph = nx.DiGraph()
        self.node_data = {}  # 存储节点数据
        self.global_config = {}  # 全局配置

    def add_node(self, node_id: str, plugin_type: str, config: Dict[str, Any] = None):
        """
        添加节点到DAG

        Args:
            node_id: 节点ID
            plugin_type: 插件类型
            config: 节点配置
        """
        if node_id in self.graph.nodes:
            raise ValueError(f"Node {node_id} already exists in DAG")

        self.graph.add_node(node_id)
        self.node_data[node_id] = {
            "plugin_type": plugin_type,
            "config": config or {}
        }

    def add_edge(self, from_node: str, to_node: str):
        """
        添加边到DAG

        Args:
            from_node: 起始节点
            to_node: 目标节点
        """
        if from_node not in self.graph.nodes:
            error_msg = f"Source node {from_node} does not exist"
            logger.error(error_msg)
            raise ValueError(error_msg)
        if to_node not in self.graph.nodes:
            error_msg = f"Target node {to_node} does not exist"
            logger.error(error_msg)
            raise ValueError(error_msg)

        self.graph.add_edge(from_node, to_node)

        # 检查是否形成环
        if not nx.is_directed_acyclic_graph(self.graph):
            self.graph.remove_edge(from_node, to_node)
            error_msg = f"Adding edge {from_node} -> {to_node} creates a cycle"
            logger.error(error_msg)
            raise ValueError(error_msg)

    def get_topological_order(self) -> List[str]:
        """
        获取拓扑排序

        Returns:
            拓扑排序的节点列表
        """
        return list(nx.topological_sort(self.graph))

    def get_dependencies(self, node_id: str) -> List[str]:
        """
        获取节点的依赖节点（前驱节点）

        Args:
            node_id: 节点ID

        Returns:
            依赖节点列表
        """
        return list(self.graph.predecessors(node_id))

    def get_ready_nodes(self, completed_nodes: Set[str]) -> List[str]:
        """
        获取可以执行的节点（所有依赖都已完成的节点）

        Args:
            completed_nodes: 已完成的节点集合

        Returns:
            可以执行的节点列表
        """
        ready_nodes = []
        for node in self.graph.nodes:
            if node in completed_nodes:
                continue

            dependencies = self.get_dependencies(node)
            if all(dep in completed_nodes for dep in dependencies):
                ready_nodes.append(node)

        return ready_nodes

    def is_complete(self, completed_nodes: Set[str]) -> bool:
        """
        检查DAG是否已完全执行

        Args:
            completed_nodes: 已完成的节点集合

        Returns:
            是否已完成
        """
        return len(completed_nodes) == len(self.graph.nodes)

    def get_node_config(self, node_id: str) -> Dict[str, Any]:
        """
        获取节点配置

        Args:
            node_id: 节点ID

        Returns:
            节点配置
        """
        return self.node_data.get(node_id, {})

    def validate(self) -> bool:
        """
        验证DAG是否有效

        Returns:
            是否有效
        """
        return nx.is_directed_acyclic_graph(self.graph)